
(* A function [f] decides a predicate [p] if it returns [true] exactly
   when [p] holds and [false] exactly when [p] does not hold. *)
predicate decide (f, p) =
    (forall x. (post (f) (x) (true)  <=> p (x)))
and (forall x. (post (f) (x) (false) <=> not (p (x))))
and (forall x. pre (f) (x))

(* Comparison operators over integers. *)
deferred val <=?  : (int * int) -> bool where decide (<=?, <=)
deferred val >=?  : (int * int) -> bool where decide (>=?, >=)
deferred val <?   : (int * int) -> bool where decide (<?, <)
deferred val >?   : (int * int) -> bool where decide (>?, >)
deferred val =?   : (int * int) -> bool
where (forall (x, y). pre ( =? ) (x, y))
 and (forall (x, y). post (=?) (x, y) (true) <=> x = y)
 and (forall (x, y). post (=?) (x, y) (false) <=> not (x = y))

(* Define a logic-level set abstraction. *)
deferred type set : *
deferred predicate %in : (int * set) ~> prop
deferred val empty : set
 where forall x. not (x %in empty)
deferred logic singleton : int ~> set
 where forall (x, y). x %in (singleton (y)) <=> (x = y)
deferred logic %U (*union*) : (set * set) ~> set
 where forall (s1, s2, x). x %in (s1 %U s2) <=> ((x %in s1) or (x %in s2))

(* Extensionality axiom. *)
predicate %=% (s1, s2) = forall x. (x %in s1) <=> (x %in s2)
axiom: forall (s1, s2). (s1 %=% s2) <=> (s1 = s2)

lemma singleton_refl: forall x. x %in singleton (x)
lemma union_assoc_ext : forall (s1, s2, s3).
(s1 %U (s2 %U s3)) %=% ((s1 %U s2) %U s3)
lemma union_assoc : forall (s1, s2, s3).
(s1 %U (s2 %U s3)) = ((s1 %U s2) %U s3)
lemma union_assoc5_ext : forall (s1, s2, s3, s4, s5).
(s1 %U s2 %U (s3 %U s4 %U s5)) %=% ((s1 %U s2 %U s3) %U s4 %U s5)
lemma union_assoc5 : forall (s1, s2, s3, s4, s5).
(s1 %U s2 %U (s3 %U s4 %U s5)) = ((s1 %U s2 %U s3) %U s4 %U s5)
lemma union_assoc4_ext : forall (s1, s2, s3, s4).
(s1 %U (s2 %U s3 %U s4)) %=% ((s1 %U s2) %U s3 %U s4)
lemma union_assoc4 : forall (s1, s2, s3, s4).
(s1 %U (s2 %U s3 %U s4)) = ((s1 %U s2) %U s3 %U s4)
lemma union_assoc4b_ext : forall (s1, s2, s3, s4).
(s1 %U s2 %U (s3 %U s4)) %=% ((s1 %U (s2 %U s3 %U s4)))
lemma union_assoc4b : forall (s1, s2, s3, s4).
(s1 %U s2 %U (s3 %U s4)) = ((s1 %U (s2 %U s3 %U s4)))
(* Too hard for ergo 0.8 !*)
axiom union_comm4_ext : forall (s1, s2, s3, s4).
((s1 %U (s2 %U s3 %U s4))) %=% ((s3 %U (s1 %U s2 %U s4)))
lemma union_comm4 : forall (s1, s2, s3, s4).
((s1 %U (s2 %U s3 %U s4))) = ((s3 %U (s1 %U s2 %U s4)))
lemma union_empty_l_ext : forall (s). (empty %U s) %=% s
lemma union_empty_l : forall (s). (empty %U s) = s
lemma union_empty_r_ext : forall (s). (s %U empty) %=% s
lemma union_empty_r : forall (s). (s %U empty) = s
lemma mem_union_ext: forall (x, s). x %in s => s %=% (singleton (x) %U s)

type color =
 | Red : color
 | Black : color

type rbtree =
 | Empty : rbtree
 | Node  : (color * rbtree * int * rbtree) -> rbtree

(* Color of tree *)
logic black(t) =
 | Empty                -> True (* empty nodes are black *)
 | Node(Black, _, _, _) -> True
 | Node(Red, _, _, _)   -> False

predicate red(t) = not (black(t))

(* Abstraction function:  convert tree to set. *)
logic toSet (l) =
 | Empty            -> empty
 | Node(_, l, x, r) -> toSet(l) %U singleton(x) %U toSet(r)

lemma not_empty_set: forall (t, x). x %in (toSet (t)) => not (t = Empty)
lemma element_in_set: forall (c, l, x, r). x %in (toSet (Node (c, l, x, r)))

(* Rep invariant *)

(* First RI:  no red node has a red child *)
(* inductive rbInv1 : rbtree ~> prop =
 | rbInv1(Empty)
 | forall (c, l, x, r).
     (((c = Red => (black(l) and black(r)))
     and rbInv1(l) and rbInv1(r)))
     => rbInv1(Node(c, l, x, r))*)

(* Second RI:  every path from the root to an empty node
 *   contains the same number of black nodes. *)
(* predicate rbInv2(t:rbtree) = True (* TODO *)*)

predicate repOk(t) = True (*rbInv1(t) and rbInv2(t)*)

(* Operations *)

let makeBlack(t:rbtree)
 where (not (t = Empty))
 returns t2 where
 ((forall (c, l, x, r). t = Node(c, l, x, r) => t2 = Node(Black, l, x, r))
  and toSet (t2) = toSet(t))
=
 match t with
 | Empty -> absurd
 | Node(_, a, y, b) -> Node(Black, a, y, b)
 end

let balance(c, l, e, r)
returns t where
(toSet(t) %=% (toSet(l) %U singleton(e) %U toSet(r)))
=
 match c with
   | Black ->
       match l with
         | Node (Red, Node (Red, a, x, b), y, c) ->
             Node(Red, Node(Black, a, x, b), y, Node(Black, c, e, r))
         | Node (Red, a, x, Node (Red, b, y, c)) ->
             Node(Red, Node(Black, a, x, b), y, Node(Black, c, e, r))
         | _ ->
             match r with
             | Node(Red, Node(Red, b, y, c), z, d) ->
                 Node(Red, Node(Black, l, e, b), y, Node(Black, c, z, d))
             | Node(Red, b, y, Node(Red, c, z, d)) ->
                 Node(Red, Node(Black, l, e, b), y, Node(Black, c, z, d))
             | _ -> Node (c, l, e, r)
             end
        end
   | _ -> Node (c, l, e, r)
  end

let rec ins(x:int, t:rbtree)
 returns t2 where (toSet(t2) %=% (singleton(x) %U toSet(t)))
=
 match t with
 | Empty -> Node(Red, Empty, x, Empty)
 | Node(c, a, y, b) ->
       if x <? y then balance(c, ins(x, a), y, b)
       else if y <? x then balance(c, a, y, ins(x, b))
       else
         assert (x = y) in
         assert (x %in toSet(t)) in t
 end

let insert (x:int, t:rbtree)
 where (repOk(t))
 returns t2 where (repOk(t2) and toSet(t2) %=% (singleton(x) %U toSet(t)))
=
 makeBlack(ins(x,t))
